import numpy as np
import matplotlib.pyplot as plt
import cv2 as cv

filename1=r'F:\AI\right03.jpg'
img1=cv.imread(filename1)
filename2=r'F:\AI\right04.jpg'
img2=cv.imread(filename2)

#sift初始化
sift=cv.xfeatures2d_SIFT.create()
k1,d1=sift.detectAndCompute(img1,None)  #显示报错，问题
k2,d2=sift.detectAndCompute(img2,None)

FLANN_INDEX_KDTREE=0
#设置FLANN超参数，FLANN最快速近邻搜索包，FLANN匹配需要两个字典（dict）index_params；search_params
index_params=dict(algorithm=FLANN_INDEX_KDTREE,tree=5)#K-D树索引超参数
search_params=dict(checks=50)#设置遍历次数，次数越高越准确但时间会长
flann=cv.FlannBasedMatcher(index_params,search_params)#初始化LannBasedMatcher匹配器
#BFMatcher是暴力搜索，FlannBasedMatcher是一种优化后的最近邻搜索方式
matches=flann.knnMatch(d1,d2,k=2)#匹配两张图特征向量

good=[]
pts1=[]
pts2=[]

#筛选较好匹配点
for i,(m,n) in enumerate(matches):
    if m.distance<0.8*n.distance:#这里是第一次匹配和次匹配的关系，α取0.8
        good.append(m)
        pts2.append(k2[m.trainIdx].pt)
        pts1.append(k1[m.queryIdx].pt)
#计算基础矩阵
pts1=np.int32(pts1)
pts2=np.int32(pts2)

F,mask=cv.findFundamentalMat(pts1,pts2,cv.FM_LMEDS)
#F是基础矩阵，mask是返回基本矩阵的值，未找到矩阵返回0，找到一个返回1，多个矩阵返回3
pts1=pts1[mask.ravel()==1]#选择有效数据
pts2=pts2[mask.ravel()==1]

def drawlines(img1,img2,lines,pts1,pts2):#绘制图像极线
    r,c=img1.shape
    img1=cv.cvtColor(img1,cv.COLOR_GRAY2BGR)
    img2=cv.cvtColor(img2,cv.COLOR_GRAY2BGR)

    for r,pt1,pt2 in zip(lines,pts1,pts2):
        color=tuple(np.random.randint(0,255,3).tolist())
        x0,y0=map(int,[0,-r[2]/r[1]])
        x1,y1=map(int,[c,-(r[2]+r[0]*c)/r[1]])

        img1=cv.line(img1,(x0,y0),(x1,y1),color,1)
        img1=cv.circle(img1,tuple(pt1),5,color,-1)
        img2 = cv.circle(img2, tuple(pt2), 5, color, -1)
    return img1,img2

#在右图找到与点对应的极线，在左图上画线
lines1=cv.computeCorrespondEpilines(pts2.reshape(-1,1,2),2,F)
lines1=lines1.reshape(-1,3)
img3,img4=drawlines(img1,img2,lines1 ,pts1,pts2)
#找到左图对应点，在右图画线
lines2 = cv.computeCorrespondEpilines(pts2.reshape(-1, 1, 2), 1, F)
lines2 = lines2.reshape(-1, 3)
img5, img6 = drawlines(img2, img1, lines2, pts2, pts1)

plt.subplot(121)
plt.imshow(img3)
plt.subplot(122)
plt.imshow(img5)
plt.show()

#用本质矩阵估计平移旋转
E,mask=cv.findEssentialMat(pts1,pts2,np.array([[4.70062852e+03, 0.00000000e+00, 1.38445962e+03],
 [0.00000000e+00, 4.72048901e+03, 2.28597865e+03],
 [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]]),method=cv.RANSAC,threshold=1,prob=0.999)
print('本质矩阵',E)

A=np.array(E)
U,sigma,VT=np.linalg.svd(A)
print('旋转',U)
print('参数',sigma)
print('平移',VT)



